package cn.iocoder.yudao.framework.websocket.core;

import cn.hutool.extra.spring.SpringUtil;
import cn.iocoder.yudao.framework.security.core.LoginUser;
import cn.iocoder.yudao.framework.websocket.config.WebSocketProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;

import javax.annotation.Resource;
import java.util.concurrent.Semaphore;
@Slf4j
public class YudaoWebSocketHandlerDecorator extends WebSocketHandlerDecorator {

    private final WebSocketProperties webSocketProperties = SpringUtil.getBean(WebSocketProperties.class);

    private final Semaphore socketSemaphore ;
    public YudaoWebSocketHandlerDecorator(WebSocketHandler delegate) {
        super(delegate);
        socketSemaphore = new Semaphore(webSocketProperties.getMaxConnection());
    }

    /**
     * websocket 连接时执行的动作
     * @param session websocket session 对象
     * @throws Exception 异常对象
     */
    @Override
    public void afterConnectionEstablished(final WebSocketSession session) throws Exception {
        super.afterConnectionEstablished(session);
        Object sessionKey = sessionKeyGen(session);
        log.info(session.getAcceptedProtocol());
        // 尝试获取信号量
        boolean semaphoreFlag = SemaphoreUtils.tryAcquire(socketSemaphore);
        if (!semaphoreFlag)
        {
            // 未获取到信号量
            log.error("\n 当前在线人数超过限制数- {}", webSocketProperties.getMaxConnection());
            WebSocketSessionHandler.sendMessageToUserByText(session, "当前在线人数超过限制数：" + webSocketProperties.getMaxConnection());
            session.close();
        }
        else
        {
            // 添加用户
            WebSocketSessionHandler.put(sessionKey.toString(), session);
            log.info("\n 建立连接 - {}", session);
            log.info("\n 当前人数 - {}", WebSocketSessionHandler.getUsers().size());
            //  WebSocketSessionHandler.sendMessageToUserByText(session, "连接成功");
        }
    }

    /**
     * websocket 关闭连接时执行的动作
     * @param session websocket session 对象
     * @param closeStatus 关闭状态对象
     * @throws Exception 异常对象
     */
    @Override
    public void afterConnectionClosed(final WebSocketSession session, CloseStatus closeStatus) throws Exception {
        super.afterConnectionClosed(session,closeStatus);
        Object sessionKey = sessionKeyGen(session);
        log.info("\n 关闭连接 - {}", session);
        // 移除用户
        WebSocketSessionHandler.remove(sessionKey.toString());
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }
    @Override
    public void handleTransportError(WebSocketSession session,  Throwable exception) throws Exception {
        super.handleTransportError(session,exception);
        if (session.isOpen())
        {
            // 关闭连接
            session.close();
        }
        String sessionId = sessionKeyGen(session).toString();
        log.info("\n 连接异常 - {}", sessionId);
        log.info("\n 异常信息 - {}", exception.getMessage());
        // 移出用户
        WebSocketSessionHandler.remove(sessionId);
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }
    @Override
    public void handleMessage(WebSocketSession session,  WebSocketMessage<?> message) throws Exception {
        super.handleMessage(session ,message);
        //TODO 这里可以预处理websocket消息
    }
    public Object sessionKeyGen(WebSocketSession webSocketSession) {

        Object obj = webSocketSession.getAttributes().get(WebSocketKeyDefine.LOGIN_USER);

        if (obj instanceof LoginUser) {
            LoginUser loginUser = (LoginUser) obj;
            // userId 作为唯一区分
            return String.valueOf(loginUser.getId());
        }

        return null;
    }
}
